# -*- coding: utf-8 -*-
"""
自动扫描和注册 MCP tools 工具的类
"""
import os
import importlib
import importlib.util
import inspect
from pathlib import Path
from typing import List, Dict, Any, Callable
from fastmcp import FastMCP

class ToolRegistry:
    """自动扫描和注册 MCP tools 工具的类"""
    
    def __init__(self, mcp_server: FastMCP):
        self.mcp = mcp_server
        self.registered_tools = []
        self.registered_resources = []
    
    def scan_directory(self, directory: str, pattern: str = "*.py") -> None:
        """
        扫描目录中的Python文件并注册工具

        Args:
            directory (str): 要扫描的目录路径
            pattern (str, optional): 文件匹配模式，默认是 "*.py"

        Returns:
            None
        """
        directory_path = Path(directory)
        
        if not directory_path.exists():
            print(f"目录不存在: {directory}")
            return
        
        # 获取所有Python文件
        python_files = directory_path.glob(pattern)
        
        for file_path in python_files:
            if file_path.name.startswith("_"):
                continue  # 跳过私有文件
                
            try:
                self._register_module(file_path)
            except Exception as e:
                print(f"注册文件失败 {file_path}: {e}")
    
    def _register_module(self, file_path: Path) -> None:
        """注册单个模块中的工具"""
        # 构建模块名
        module_name = file_path.stem
        
        # 动态导入模块
        spec = importlib.util.spec_from_file_location(module_name, file_path)
        if spec is None or spec.loader is None:
            return
            
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        
        # 扫描模块中的所有函数
        for name, obj in inspect.getmembers(module, inspect.isfunction):
            self._register_function(obj)
    
    def _register_function(self, func: Callable) -> None:
        """注册单个函数"""
        # 检查是否是工具函数
        if hasattr(func, '_mcp_tool'):
            self.mcp.tool()(func) # 将扫描到的 MCP tools 注册到 MCP Server 中
            self.registered_tools.append(func.__name__)
            print(f"注册工具: {func.__name__}")
        
        # 检查是否是资源函数
        if hasattr(func, '_mcp_resource'):
            uri = func._mcp_resource_uri
            self.mcp.resource(uri)(func) # 将扫描到的 MCP resources 注册到 MCP Server 中
            self.registered_resources.append((func.__name__, uri))
            print(f"注册资源: {func.__name__} -> {uri}")

# 创建装饰器函数
def mcp_tool():
    """标记函数为MCP工具的装饰器"""
    def decorator(func):
        func._mcp_tool = True
        return func
    return decorator

def mcp_resource(uri: str):
    """标记函数为MCP资源的装饰器"""
    def decorator(func):
        func._mcp_resource = True
        func._mcp_resource_uri = uri
        return func
    return decorator